# src/text_preprocessor.py
import os
import re
import json
import math
from multiprocessing import Pool, cpu_count
import tqdm  # 确保tqdm已安装

# 假设您有这样的utils和config导入
from src.utils import setup_logger  # , get_progress_bar (tqdm can be used directly)
from config import MIN_SENTENCE_LENGTH  # , OTHER_CONFIGS (if needed)

logger = setup_logger(__name__, "text_preprocessor.log")


# 辅助函数：从单个文件提取句子 (这是您的原始逻辑)
def extract_sentences_from_file_worker(args):
    """
    Worker function to extract sentences from a single file.
    Args:
        filepath (str): Path to the text file.
        min_length (int): Minimum length of a sentence.
    Returns:
        list: A list of sentences extracted from the file.
    """
    filepath, min_length = args
    logger.debug(f"正在从文件提取句子: {filepath}")
    sentences = []
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        # 您可能需要更复杂的句子切分逻辑
        # 这里使用一个基础的正则切分示例
        content = content.replace(' ','').replace('\n','')
        raw_sentences = re.split(r'[.。！!?？\n]+', content)
        for s in raw_sentences:
            s_cleaned = s.strip()
            if s_cleaned and len(s_cleaned) >= min_length:
                try:
                    float(s_cleaned.replace(',',''))
                    continue
                except ValueError:
                    pass
                sentences.append(s_cleaned)

    except Exception as e:
        logger.error(f"读取或处理文件 {filepath} 时出错: {e}")
    # logger.info(f"从 {filepath} 提取了 {len(sentences)} 条句子。") # 日志可能过于频繁
    return sentences


def process_files_chunk(file_chunk_args):
    """
    Processes a chunk of files in parallel to extract sentences.
    Helper for process_all_text_files.
    Args:
        file_chunk_args (tuple): Contains (list_of_filepaths, min_sentence_length, num_processes_for_chunk)
    Returns:
        list: Aggregated list of sentences from this chunk of files.
    """
    files_in_chunk, min_length, num_processes = file_chunk_args
    all_sentences_in_chunk = []

    # 为当前文件块设置进程池
    # 注意：如果文件块很小，为每个块创建和销毁进程池可能会有开销。
    # 对于大量小文件，另一种策略是有一个长期存在的池，并将所有文件分派给它。
    # 但按块处理有助于控制内存和保存的粒度。
    # 我们限制子进程池的大小，例如，不超过 cpu_count() // 2 或一个合理的值，
    # 以避免创建过多的进程，特别是如果 process_all_text_files 本身就是由一个更大的并行框架调用的。
    # 这里我们假设 process_all_text_files 是顶层调用者。

    actual_num_processes = min(num_processes, len(files_in_chunk)) if files_in_chunk else 1

    logger.info(f"正在使用 {actual_num_processes} 个进程处理包含 {len(files_in_chunk)} 个文件的块...")

    # Prepare arguments for worker function
    tasks = [(filepath, min_length) for filepath in files_in_chunk]

    with Pool(processes=actual_num_processes) as pool:
        # 使用 imap 和 tqdm 来获得进度条
        # results_iter = pool.imap(extract_sentences_from_file_worker, tasks)
        # for sentences_from_file in tqdm.tqdm(results_iter, total=len(tasks), desc="处理文件块内的文件"):
        #     if sentences_from_file:
        #         all_sentences_in_chunk.extend(sentences_from_file)
        #
        #  不使用子tqdm，避免日志混乱，让外部tqdm处理块的进度
        results = pool.map(extract_sentences_from_file_worker, tasks)
        for sentences_from_file in results:
            if sentences_from_file:
                all_sentences_in_chunk.extend(sentences_from_file)

    logger.info(f"文件块处理完成，提取了 {len(all_sentences_in_chunk)} 条句子。")
    return all_sentences_in_chunk


def process_all_text_files(all_report_filepaths, min_sentence_length, sentence_json_dir,
                           files_per_chunk=100, num_processes=None):
    """
    处理所有文本文件，提取句子，并将它们分块保存到JSON文件中。
    Args:
        all_report_filepaths (list): 所有报告文件的路径列表。
        min_sentence_length (int): 句子的最小长度。
        sentence_json_dir (str): 保存分块句子JSON文件的目录。
        files_per_chunk (int): 每个JSON块包含多少个文件的句子。
        num_processes (int, optional): 用于并行处理文件的进程数。默认为 cpu_count()。
    """
    if not all_report_filepaths:
        logger.warning("没有提供报告文件路径，跳过句子提取。")
        return

    if num_processes is None:
        num_processes = cpu_count()
    logger.info(f"将使用最多 {num_processes} 个进程进行文件级并行处理。")

    os.makedirs(sentence_json_dir, exist_ok=True)  # 确保输出目录存在

    total_files = len(all_report_filepaths)
    num_chunks = math.ceil(total_files / files_per_chunk)

    logger.info(f"总共 {total_files} 个文件, 将分为 {num_chunks} 块，每块最多 {files_per_chunk} 个文件。")

    processed_sentence_count_total = 0
    output_file_paths = []

    for i in tqdm.tqdm(range(num_chunks), desc="处理文件块"):
        chunk_start_index = i * files_per_chunk
        chunk_end_index = min((i + 1) * files_per_chunk, total_files)
        current_file_chunk_paths = all_report_filepaths[chunk_start_index:chunk_end_index]

        if not current_file_chunk_paths:
            logger.info(f"块 {i + 1}/{num_chunks} 为空，跳过。")
            continue

        logger.info(f"开始处理块 {i + 1}/{num_chunks}, 包含 {len(current_file_chunk_paths)} 个文件。")

        # 将文件列表和参数传递给并行处理函数
        # 注意：这里我们将整个文件块交给 process_files_chunk，它内部会创建自己的进程池。
        # 如果希望全局控制进程数，这里的 num_processes 可以直接传递给 process_files_chunk
        # 并在那里直接使用，而不是让它再次计算。
        # 为简单起见，这里让 process_files_chunk 自行管理其进程。
        sentences_in_chunk = process_files_chunk((current_file_chunk_paths, min_sentence_length, num_processes))

        if sentences_in_chunk:
            output_filename = f"sentences_chunk_{i + 1}.json"
            output_filepath = os.path.join(sentence_json_dir, output_filename)

            try:
                with open(output_filepath, 'w', encoding='utf-8') as f:
                    json.dump(sentences_in_chunk, f, ensure_ascii=False, indent=4)
                logger.info(f"块 {i + 1}/{num_chunks} 的 {len(sentences_in_chunk)} 条句子已保存到: {output_filepath}")
                processed_sentence_count_total += len(sentences_in_chunk)
                output_file_paths.append(output_filepath)
            except Exception as e:
                logger.error(f"无法保存块 {i + 1}/{num_chunks} 的句子到 {output_filepath}: {e}")
        else:
            logger.info(f"块 {i + 1}/{num_chunks} 未提取到句子。")

    logger.info(f"所有文件块处理完成。总共提取并保存了 {processed_sentence_count_total} 条句子。")
    logger.info(f"句子JSON文件保存在目录: {sentence_json_dir}")
    # 这个函数现在不返回句子列表，调用者需要从JSON文件中读取
    # 可以选择返回生成的文件路径列表
    return output_file_paths  # 或者不返回任何东西，让调用者扫描目录


# 如果需要，可以在这里添加独立的测试 __main__
if __name__ == '__main__':
    logger.info("独立测试 src/text_preprocessor.py...")

    # 创建一些假的测试文件和配置
    TEST_DATA_DIR = "temp_test_data"
    TEST_REPORTS_DIR = os.path.join(TEST_DATA_DIR, "reports")
    TEST_SENTENCE_JSON_DIR = os.path.join(TEST_DATA_DIR, "sentence_chunks_test")

    os.makedirs(TEST_REPORTS_DIR, exist_ok=True)
    os.makedirs(TEST_SENTENCE_JSON_DIR, exist_ok=True)

    # 创建示例报告文件
    sample_files = []
    for i in range(5):  # 创建5个示例文件
        fn = os.path.join(TEST_REPORTS_DIR, f"report_{i + 1}.txt")
        with open(fn, 'w', encoding='utf-8') as f:
            f.write(f"这是报告{i + 1}的第一句话。这是第二句话。\n文件的第三句话也很重要！")
            if i % 2 == 0:
                f.write("只有偶数报告有第四句话。")
        sample_files.append(fn)

    # 配置参数
    min_len = 8
    chunk_size = 2  # 每块处理2个文件

    logger.info(f"开始使用 {len(sample_files)} 个示例文件进行测试...")
    process_all_text_files(sample_files, min_len, TEST_SENTENCE_JSON_DIR,
                           files_per_chunk=chunk_size, num_processes=2)

    logger.info(f"测试完成。检查 '{TEST_SENTENCE_JSON_DIR}' 中的输出。")

    # 清理测试文件（可选）
    # import shutil
    # shutil.rmtree(TEST_DATA_DIR)
    # logger.info(f"已清理测试目录: {TEST_DATA_DIR}")